from llama_index.core.agent import ReActAgent
from llama_index.core.tools import FunctionTool
from llama_index.llms.openllm import OpenLLM
from llama_index.llms.openai import OpenAI
import os

def add(x: int, y: int) -> int:
    """Adds two integers together."""
    return x + y


def mystery(x: int, y: int) -> int:
    """Mystery function that operates on top of two numbers."""
    return (x + y) * (x + y)


add_tool = FunctionTool.from_defaults(fn=add)
mystery_tool = FunctionTool.from_defaults(fn=mystery)

# llm = OpenAI(model="gpt-3.5-turbo")

llm = OpenLLM(
    model="deepseek-r1:7b", api_base="http://127.0.0.1:11434/v1", api_key="Empty"
)

# llm = OpenLLM(
#     model="deepseek-ai/DeepSeek-R1-Distill-Qwen-7B", api_base="https://api.siliconflow.cn/v1", api_key=os.environ.get("SILICON_API_KEY")
# )

# 使用react agent
agent = ReActAgent.from_tools([add_tool, mystery_tool], llm=llm, verbose=True)
response = agent.chat("Tell me 2 add 5 mystery 3")
print(str(response))
